from PyQt5.QtCore import QThread, pyqtSignal
import os
import numpy as np
from astropy.io import fits
import astroalign as aa

class BatchAlignThread(QThread):
    """批量图像对齐线程"""
    progress_signal = pyqtSignal(int, str)  # 发送当前处理的索引和文件名
    complete_signal = pyqtSignal()  # 对齐完成信号
    error_signal = pyqtSignal(str)  # 错误信号
    
    def __init__(self, new_folder, old_folder, aligned_new_folder, new_files, old_files):
        super().__init__()
        self.new_folder = new_folder
        self.old_folder = old_folder
        self.aligned_new_folder = aligned_new_folder
        self.new_files = new_files
        self.old_files = old_files
        self.is_canceled = False
    
    def run(self):
        try:
            # 遍历新图文件
            for idx, new_file in enumerate(self.new_files):
                if self.is_canceled:
                    break
                
                try:
                    # 发送进度信号
                    self.progress_signal.emit(idx, new_file)
                    
                    # 提取新图文件名（不含扩展名）
                    base_name = os.path.splitext(new_file)[0]
                    # 查找旧图文件夹中匹配的文件
                    matching_old_files = [f for f in self.old_files if base_name in os.path.splitext(f)[0]]
                    if not matching_old_files:
                        print(f"未找到与 {new_file} 匹配的旧图，跳过此图像。")
                        continue
                    old_file = matching_old_files[0]

                    # 读取图像
                    new_img_path = os.path.join(self.new_folder, new_file)
                    old_img_path = os.path.join(self.old_folder, old_file)
                    
                    try:
                        # 直接使用fits.getdata读取数据而不处理头信息
                        new_img = fits.getdata(new_img_path)
                        old_img = fits.getdata(old_img_path)
                        
                        # 检查图像维度
                        if len(new_img.shape) > 2:
                            new_img = new_img[0]  # 取第一个通道
                        if len(old_img.shape) > 2:
                            old_img = old_img[0]  # 取第一个通道
                        
                        # 转换为浮点型
                        new_img = new_img.astype(np.float32)
                        old_img = old_img.astype(np.float32)
                        
                        # 使用astroalign对齐图像
                        aligned_new_img, footprint = aa.register(new_img, old_img)
                        
                        # 保存对齐后的图像
                        aligned_new_img_path = os.path.join(self.aligned_new_folder, new_file)
                        
                        # 创建新的FITS文件
                        hdu_new = fits.PrimaryHDU(aligned_new_img.astype(np.float32))
                        hdu_new.header['OBJECT'] = 'Aligned New Image'
                        hdu_new.writeto(aligned_new_img_path, overwrite=True)
                        
                        print(f"图像 {new_file} 已对齐并保存到 {aligned_new_img_path}")
                    except Exception as e:
                        print(f"处理图像 {new_file} 和 {old_file} 时出现错误: {e}，跳过此图像。")
                        # 如果对齐失败，跳过该图像，不复制原始图像
                        continue
                
                except Exception as e:
                    print(f"处理图像 {new_file} 时出错: {e}，跳过此图像。")
            
            # 发送完成信号
            self.complete_signal.emit()
        
        except Exception as e:
            # 如果有未捕获的异常，发送错误信号
            self.error_signal.emit(str(e))
    
    def cancel(self):
        """取消对齐操作"""
        self.is_canceled = True